// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/future/mutex.h>

#include <mutex>

#ifndef _WIN32

#include <pthread.h>
#include <atomic>

#endif

#include <nebula/version.h>
#include <turbo/log/logging.h>

namespace nebula::util {

    struct Mutex::Impl {
        std::mutex mutex_;
    };

    Mutex::Guard::Guard(Mutex *locked)
            : locked_(locked, [](Mutex *locked) {
        DKCHECK(!locked->impl_->mutex_.try_lock());
        locked->impl_->mutex_.unlock();
    }) {
    }

    Mutex::Guard Mutex::TryLock() {
        DKCHECK(impl_ != nullptr);
        if (impl_->mutex_.try_lock()) {
            return Guard{this};
        }
        return Guard{};
    }

    Mutex::Guard Mutex::Lock() {
        DKCHECK(impl_ != nullptr);
        impl_->mutex_.lock();
        return Guard{this};
    }

    Mutex::Mutex() : impl_(new Impl, [](Impl *impl) { delete impl; }) {}

#ifndef _WIN32
    namespace {

        struct AfterForkState {
            // A global instance that will also register the atfork handler when
            // constructed.
            static AfterForkState instance;

            // The mutex may be used at shutdown, so make it eternal.
            // The leak (only in child processes) is a small price to pay for robustness.
            Mutex *mutex = nullptr;

        private:
            AfterForkState() {
                pthread_atfork(/*prepare=*/nullptr, /*parent=*/nullptr, /*child=*/&AfterFork);
            }

            static void AfterFork() { instance.mutex = new Mutex; }
        };

        AfterForkState AfterForkState::instance;
    }  // namespace

    Mutex *GlobalForkSafeMutex() { return AfterForkState::instance.mutex; }

#endif  // _WIN32

}  // namespace nebula::util

